#include "sbox.h"
#include "one_round.h"
#include "types.h"
#include "matrix.h"
#include "field.h"
#include "tables.h"
#include "pht.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

extern byte gf_256_169_mul(byte,byte);

void guess_h_func(byte* left_part, byte* des, byte*sbox)
{
	int i;
	byte tmp[BLOCK_SIZE/2];
	memcpy(tmp, left_part, sizeof(tmp));
	for(i = 0; i < BLOCK_SIZE/2; ++i)
	{
		tmp[i] = sbox[tmp[i]];
	}
	
	byte matrix_tmp[4];

	byte_matrix_mul(mds, 4, 4, tmp, 4, 1, matrix_tmp, sizeof(matrix_tmp), element_add, gf_256_169_mul);
	memcpy(tmp, matrix_tmp ,4);

	byte_matrix_mul(mds, 4, 4, tmp + 4, 4, 1, matrix_tmp, sizeof(matrix_tmp), element_add, gf_256_169_mul);
	memcpy(tmp + 4, matrix_tmp, 4);

	memcpy(des, tmp, BLOCK_SIZE/2);
}

void guess_one_round(byte* left_part, 
					byte* right_part, 
					word* sub_key,
					byte* des,
					byte* sbox)
{
	int i;
	byte tmp[BLOCK_SIZE/2];

	/*
	memcpy(tmp, left_part, sizeof(tmp));
	for(i = 0; i < BLOCK_SIZE/2; ++i)
	{
		tmp[i] = sbox[tmp[i]];
	}
	
	byte matrix_tmp[4];

	byte_matrix_mul(mds, 4, 4, tmp, 4, 1, matrix_tmp, sizeof(matrix_tmp), element_add, gf_256_169_mul);
	word l0 = join_word(matrix_tmp);

	byte_matrix_mul(mds, 4, 4, tmp + 4, 4, 1, matrix_tmp, sizeof(matrix_tmp), element_add, gf_256_169_mul);
	word l1 = join_word(matrix_tmp);
	*/

	guess_h_func(left_part, tmp, sbox);
	word l0 = join_word(tmp);
	word l1 = join_word(tmp+sizeof(word));

	pht(l0, l1, &l0, &l1);

	printf("l0:%8x l1:%8x\n", l0, l1);

	l0 = l0 + sub_key[0];
	l1 = l1 + sub_key[1];

	split_word(l0, tmp, 4);
	split_word(l1, tmp+4, 4);

	for(i = 0; i < sizeof(tmp); ++i)
	{
		tmp[i] ^= right_part[i];
	}
	memcpy(des, tmp, sizeof(tmp));
}

byte add_in_diff[0x100][0x100][0x100];
int add_in_diff_size[0x100][0x100];

const char* file_name = "od_ad_id.rlt";
void save_to_file()
{
	FILE* save_file = fopen(file_name, "wb+");
	if(save_file == NULL)
	{
		printf("can't open file %s\n", file_name);
		exit(1);
	}

	int try_od,try_ad;
	for(try_od = 0; try_od <= 0xff; ++try_od)
	{
		for(try_ad = 0; try_ad <= 0xff; ++try_ad)
		{
			byte out_diff = try_od;
			int try_in_diff,try_adder;
			add_in_diff_size[try_od][try_ad] = 0;
			for(try_in_diff = 0; try_in_diff <= 0xff; ++try_in_diff)
			{
				int found = 0;
				int try_in;
				for(try_in = 0; try_in <= 0xff; ++try_in)
				{
					int try_in_e = try_in ^ try_in_diff;
					for(try_adder = 0; try_adder <= 0xff; ++try_adder)
					{
						int try_adder_e = try_adder + try_ad;
						byte out = try_in + try_adder;
						byte out_e = try_in_e + try_adder_e;
						byte od = out ^ out_e;
						if(od == out_diff)
						{
							found = 1;
							break;
						}
					}
					if(found)
						break;
				}
				if(found)
				{
					int size = add_in_diff_size[try_od][try_ad];
					add_in_diff[try_od][try_ad][size] = try_in_diff;
					++add_in_diff_size[try_od][try_ad];
				}
			}
		}
		printf("try_od:%d\n", try_od);
	}

	for(try_od = 0; try_od <= 0xff; ++try_od)
	{
		for(try_ad = 0; try_ad <= 0xff; ++try_ad)
		{
			int size = add_in_diff_size[try_od][try_ad];
			int write_size = fwrite(&size, sizeof(size), 1, save_file);
			if(write_size != 1)
			{
				printf("fwrite error!\n");
				exit(1);
			}
		}
	}

	for(try_od = 0; try_od <= 0xff; ++try_od)
	{
		for(try_ad = 0; try_ad <= 0xff; ++try_ad)
		{
			int size = add_in_diff_size[try_od][try_ad];
			int j;
			for( j = 0; j < size; ++j)
			{
				byte value = add_in_diff[try_od][try_ad][j];
				int write_size = fwrite(&value, sizeof(value), 1, save_file);
				if(write_size != 1)
				{
					printf("fwrite error!\n");
					exit(1);
				}
			}
		}
	}

	fclose(save_file);
}

void read_from_file()
{
	FILE* save_file = fopen(file_name, "rb");
	if(save_file == NULL)
	{
		printf("open file %s to read error\n", file_name);
		exit(1);
	}

	int try_od, try_ad;
	for(try_od = 0; try_od <= 0xff; ++try_od)
	{
		for(try_ad = 0; try_ad <= 0xff; ++try_ad)
		{
			int read_count = fread(&add_in_diff_size[try_od][try_ad], sizeof(int), 1, save_file);
			if(read_count != 1)
			{
				printf("fread error!\n");
				exit(1);
			}
		}
	}

	for(try_od = 0; try_od <= 0xff; ++try_od)
	{
		for(try_ad = 0; try_ad <= 0xff; ++try_ad)
		{
			int size = add_in_diff_size[try_od][try_ad];
			int j;
			for( j = 0; j < size; ++j)
			{
				int read_count = fread(&add_in_diff[try_od][try_ad][j], sizeof(byte), 1, save_file);
				if(read_count != 1)
				{
					printf("fread error!\n");
					exit(1);
				}
			}
		}
	}

	fclose(save_file);
}

int main()
{
	int error = 0x01;
	byte plaintext[BLOCK_SIZE] = {
		0
	};
	byte p_e[BLOCK_SIZE] = {
		plaintext[0] ^ error,0
	};

	byte result[BLOCK_SIZE];
	byte result_e[BLOCK_SIZE];
	word sub_key[2] = {
		0x37983212,
		0xa32fed98,
	};
	
	for(error = 0x01; error <= 0xff; ++error)
	{
	printf("Error value:%.2x\n", error);
	p_e[0] = plaintext[0] ^ error;


	one_round(plaintext, plaintext + BLOCK_SIZE/2, sub_key, result);
	int i;
	for( i = 0; i < BLOCK_SIZE/2; ++i)
	{
		printf("%.2x ", result[i]);
	}
	printf("\n");

	one_round(p_e, p_e + BLOCK_SIZE/2, sub_key, result_e);
	for( i = 0; i < BLOCK_SIZE/2; ++i)
	{
		printf("%.2x ", result_e[i]);
	}
	printf("\n");

	byte guess_sbox[256] = {
		0
	};
	
	read_from_file();

	int cand[0xffff];
	int size = 0;
	byte guess_result[BLOCK_SIZE];
	byte guess_result_e[BLOCK_SIZE];
	int zero_guess, e_guess;
	for(zero_guess = 0; zero_guess <= 0xff; ++zero_guess)
	{
		guess_sbox[0] = zero_guess;
		for( e_guess = 0; e_guess <= 0xff; ++e_guess)
		{
			guess_sbox[error] = e_guess;
			guess_h_func(plaintext, guess_result, guess_sbox);
			guess_h_func(p_e, guess_result_e, guess_sbox);

			word l0 = join_word(guess_result);
			word l0_e = join_word(guess_result_e);
			word l1 = join_word(guess_result + 4);
			word l1_e = join_word(guess_result_e + 4);

//			pht(l0, l1, &l0, &l1);
//			pht(l0_e, l1_e, &l0_e, &l1_e);

			split_word(l0, guess_result, 4);
			split_word(l0_e, guess_result_e, 4);
			split_word(l1, guess_result+4, 4);
			split_word(l1_e, guess_result_e+4, 4);

			
			word l0_sub = l0_e - l0;
			byte ads[8] = {
				0
			};
			split_word(l0_sub, ads + 4, 4);
			

			int flag = 1;
			for(i = 0; i < BLOCK_SIZE/2; ++i)
			{
				int found = 0;
				byte od = result[i] ^ result_e[i];
				byte guess_id = guess_result[i] ^ guess_result_e[i];
				int ad;
				for(ad = ads[i] ; ad <= ads[i] + 1; ++ad)
				{
					//if( (i % 4) == 0 && ad == 1)
					//	continue;
					if( (i % 4) == 0 && ad != ads[i])
						continue;
					int j;
					for( j = 0; j < add_in_diff_size[od][ad]; ++j)
					{
						if(guess_id == add_in_diff[od][ad][j])
						{
							found = 1;
							break;
						}
					}
				}
				if(!found)
				{
					flag = 0;
					break;
				}
			}
			if(flag)
			{
				cand[size] = (e_guess<<8) | zero_guess;
				size++;
			}
		}
	}

	printf("size:%d\n", size);
/*	for(i = 0; i < size; ++i)
	{
		int value = cand[i];
		int zero,e;
		zero = value & 0xff;
		e = (value >> 8) & 0xff;
		printf("%d:%d  %d:%d\n",plaintext[0], zero, p_e[0], e);
	}*/
	printf("real value;%d:%d %d:%d\n", plaintext[0], sbox[plaintext[0]], p_e[0], sbox[p_e[0]]);

	int in = 0;
	for( i = 0; i < size; ++i)
	{
		int value = cand[i];
		int zero,e;
		zero = value & 0xff;
		e = (value >> 8) & 0xff;
		if( zero == sbox[plaintext[0]] && sbox[p_e[0]] == e)
		{
			in = 1;
			break;
		}
	}

	if(in)
	{
		printf("Real value is in set!\n");
	}
	else
	{
		printf("Real value missed!!!!!!!!\n");
	}


	}

	return 0;
}
